package com.snail.tools;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.query_dsl.BoolQuery;
import co.elastic.clients.elasticsearch._types.query_dsl.RangeQuery;
import co.elastic.clients.elasticsearch.cat.IndicesRequest;
import co.elastic.clients.elasticsearch.cat.IndicesResponse;
import co.elastic.clients.elasticsearch.cat.indices.IndicesRecord;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.indices.GetMappingRequest;
import co.elastic.clients.json.JsonData;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.snail.api.ESSearchFeignClient;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.tool.execution.ToolExecutionException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.time.Instant;
import java.util.*;

/**
 * @author cici
 * @description es服务注册成toolcallback
 * @date 2025/8/18 10:33
 */
@Slf4j
@Component
public class ESsearchTool implements ToolCallback {



    private final ESSearchFeignClient feignClient;
    @Autowired
    public ESsearchTool(ESSearchFeignClient feignClient) {
        this.feignClient = feignClient;
    }
    @Autowired
    private ElasticsearchClient elasticsearchClient;

    private final ObjectMapper mapper = new ObjectMapper();
    // 实现与官方MCP服务器相同的接口规范
    @Override
    public ToolDefinition getToolDefinition() {
        // JSON Schema 必须是字符串；这里用 Java 文本块，需要 JDK 17+
        return ToolDefinition.builder()
                .name("elasticsearch_mcp")
                .description("直接调用 Elasticsearch 搜索索引")
                .inputSchema("""
                    {
                      "type": "object",
                      "properties": {
                        "entity": {
                          "type": "string",
                          "enum": ["product","user","work","collection","order","review"],
                          "description": "搜索目标类型"
                        },
                        "keyword": {
                          "type": "string",
                          "description": "搜索关键字"
                        },
                        "page": {
                          "type": "integer",
                          "default": 1
                        },
                        "size": {
                          "type": "integer",
                          "default": 10,
                          "maximum": 100
                        }
                      },
                      "required": ["entity","keyword"]
                    }
                    """)
                .build();
    }
    @Override
    public String call(String toolInput) {
        log.info("接收到搜索请求: " + toolInput);
        try {
            JsonNode in = mapper.readTree(toolInput == null ? "{}" : toolInput);
            String entity = in.path("entity").asText();
            String keyword = in.path("keyword").asText();
            int page = Math.max(in.path("page").asInt(1), 1);
            int size = Math.min(Math.max(in.path("size").asInt(10), 1), 100);
            String startDate = in.path("startDate").asText(null);
            String endDate = in.path("endDate").asText(null);
            Integer status = in.has("status") ? in.path("status").asInt() : null;

            String index;
            List<String> searchFields;
            switch (entity) {
                case "product":
                    index = "biz_product";
                    searchFields = Arrays.asList("product_name","description");
                    break;
                case "user":
                    index = "sys_user";
                    searchFields = Arrays.asList("username","nickname");
                    break;
                case "work":
                    index = "biz_works";
                    searchFields = Arrays.asList("works_name","description");
                    break;
                case "collection":
                    index = "biz_collection";
                    searchFields = Arrays.asList("works_name","name","description");
                    break;
                case "order":
                    index = "biz_order";
                    searchFields = Arrays.asList("product_name","description");
                    break;
                case "review":
                    index = "biz_product_review";
                    searchFields = Arrays.asList("content","description");
                    break;
                default: throw new IllegalArgumentException("未知实体类型: " + entity);
            }

            // 构建动态查询
            SearchResponse<JsonNode> response = elasticsearchClient.search(s -> s
                            .index(index)
                            .from((page - 1) * size)
                            .size(size)
                            .query(q -> q
                                    .bool(b -> {
                                        BoolQuery.Builder builder = b;
                                        if (keyword != null && !keyword.isEmpty()) {
                                            builder.must(m -> m
                                                    .multiMatch(mm -> mm
                                                            .query(keyword)
                                                            .fields(searchFields)
                                                    )
                                            );
                                        }
                                        if (status != null) {
                                            builder.filter(f -> f
                                                    .term(t -> t
                                                            .field("status")
                                                            .value(status)
                                                    )
                                            );
                                        }
                                        if (startDate != null || endDate != null) {
                                            builder.filter(f -> f
                                                    .range(r -> {
                                                        RangeQuery.Builder rb = r.field("create_time");
                                                        if (startDate != null) rb.gte(JsonData.of(startDate));
                                                        if (endDate != null) rb.lte(JsonData.of(endDate));
                                                        return rb;
                                                    })
                                            );
                                        }
                                        return builder;
                                    })
                            ),
                    JsonNode.class
            );

            ArrayNode hits = mapper.createArrayNode();
            response.hits().hits().forEach(hit -> hits.add(hit.source()));
            ObjectNode result = mapper.createObjectNode();
            result.put("total", response.hits().total().value());
            result.set("hits", hits);

            return mapper.writeValueAsString(result);

        } catch (Exception e) {
            throw new ToolExecutionException(getToolDefinition(), e);
        }

}
}




//    @Override
//    public String call(String toolInput) {
//        try {
//            JsonNode input = objectMapper.readTree(toolInput == null ? "{}" : toolInput);
//            String operation = input.path("operation").asText();
//
//            switch (operation) {
//                case "list_indices":
//                    return listIndices();
//                case "get_mappings":
//                    String indexForMapping = input.path("index").asText();
//                    if (indexForMapping.isEmpty()) {
//                        throw new IllegalArgumentException("index参数是get_mappings操作必需的");
//                    }
//                    return getMappings(indexForMapping);
//                case "search":
//                    return performSearch(input);
//                case "esql":
//                    String esqlQuery = input.path("esql_query").asText();
//                    if (esqlQuery.isEmpty()) {
//                        throw new IllegalArgumentException("esql_query参数是esql操作必需的");
//                    }
//                    String format = input.path("format").asText("json");
//                    return executeEsql(esqlQuery, format);
//                case "get_shards":
//                    return getShards(input.path("index").asText());
//                default:
//                    throw new IllegalArgumentException("不支持的操作类型: " + operation);
//            }
//        } catch (Exception e) {
//            // 返回与官方MCP服务器兼容的错误格式
//            return createErrorResponse(e.getMessage(), e.getClass().getSimpleName());
//        }
//    }
//
//    private String listIndices() throws Exception {
//        try {
//            IndicesRequest request = IndicesRequest.of(builder ->
//                    builder.format("json")
//                            .h("index", "health", "status", "pri", "rep", "docs.count", "docs.deleted", "store.size", "pri.store.size")
//            );
//
//            IndicesResponse response = elasticsearchClient.cat().indices(request);
//
//            List<Map<String, Object>> indices = new ArrayList<>();
//            for (IndicesRecord record : response.valueBody()) {
//                Map<String, Object> indexInfo = new HashMap<>();
//                indexInfo.put("index", record.index());
//                indexInfo.put("health", record.health());
//                indexInfo.put("status", record.status());
//                indexInfo.put("pri", record.pri());
//                indexInfo.put("rep", record.rep());
//                indexInfo.put("docs_count", record.docsCount());
//                indexInfo.put("docs_deleted", record.docsDeleted());
//                indexInfo.put("store_size", record.storeSize());
//                indexInfo.put("pri_store_size", record.priStoreSize());
//                indices.add(indexInfo);
//            }
//
//            Map<String, Object> result = new HashMap<>();
//            result.put("operation", "list_indices");
//            result.put("success", true);
//            result.put("timestamp", Instant.now().toString());
//            result.put("total_indices", indices.size());
//            result.put("indices", indices);
//
//            return objectMapper.writeValueAsString(result);
//
//        } catch (Exception e) {
//            throw new RuntimeException("获取索引列表失败: " + e.getMessage(), e);
//        }
//    }
//
//    private String getMappings(String indexName) throws Exception {
//        try {
//            GetMappingRequest request = GetMappingRequest.of(builder ->
//                    builder.index(indexName)
//            );
//
//            GetMappingResponse response = elasticsearchClient.indices().getMapping(request);
//
//            Map<String, Object> result = new HashMap<>();
//            result.put("operation", "get_mappings");
//            result.put("success", true);
//            result.put("timestamp", Instant.now().toString());
//            result.put("index", indexName);
//            result.put("mappings", response.result());
//
//            return objectMapper.writeValueAsString(result);
//
//        } catch (Exception e) {
//            throw new RuntimeException("获取索引映射失败 [" + indexName + "]: " + e.getMessage(), e);
//        }
//    }
//
//    private String performSearch(JsonNode inputNode) throws Exception {
//        try {
//            String indexName = inputNode.path("index").asText();
//            if (indexName.isEmpty()) {
//                throw new IllegalArgumentException("index参数是search操作必需的");
//            }
//
//            int size = Math.min(inputNode.path("size").asInt(10), 100); // 限制最大100
//            int from = Math.max(inputNode.path("from").asInt(0), 0);
//            JsonNode queryNode = inputNode.path("query");
//
//            SearchRequest.Builder searchBuilder = new SearchRequest.Builder()
//                    .index(indexName)
//                    .size(size)
//                    .from(from);
//
//            if (!queryNode.isMissingNode() && !queryNode.isNull()) {
//                // 如果提供了查询DSL，使用它
//                String queryJson = objectMapper.writeValueAsString(queryNode);
//                Query query = Query.of(q -> q.withJson(
//                        objectMapper.readTree(queryJson)
//                ));
//                searchBuilder.query(query);
//            } else {
//                // 默认查询所有
//                searchBuilder.query(Query.of(q -> q.matchAll(m -> m)));
//            }
//
//            SearchResponse<ObjectNode> response = elasticsearchClient.search(
//                    searchBuilder.build(),
//                    ObjectNode.class
//            );
//
//            List<Map<String, Object>> hits = new ArrayList<>();
//            for (Hit<ObjectNode> hit : response.hits().hits()) {
//                Map<String, Object> hitMap = new HashMap<>();
//                hitMap.put("_id", hit.id());
//                hitMap.put("_index", hit.index());
//                hitMap.put("_score", hit.score());
//                hitMap.put("_source", hit.source());
//
//                // 添加高亮显示结果（如果有）
//                if (hit.highlight() != null && !hit.highlight().isEmpty()) {
//                    hitMap.put("highlight", hit.highlight());
//                }
//
//                hits.add(hitMap);
//            }
//
//            Map<String, Object> result = new HashMap<>();
//            result.put("operation", "search");
//            result.put("success", true);
//            result.put("timestamp", Instant.now().toString());
//            result.put("index", indexName);
//            result.put("took", response.took());
//            result.put("timed_out", response.timedOut());
//            result.put("total_hits", response.hits().total().value());
//            result.put("total_relation", response.hits().total().relation().jsonValue());
//            result.put("max_score", response.hits().maxScore());
//            result.put("hits", hits);
//            result.put("query_params", Map.of("size", size, "from", from));
//
//            return objectMapper.writeValueAsString(result);
//
//        } catch (Exception e) {
//            throw new RuntimeException("执行搜索失败: " + e.getMessage(), e);
//        }
//    }
//
//    private String executeEsql(String esqlQuery, String format) throws Exception {
//        try {
//            EsqlQueryRequest request = EsqlQueryRequest.of(builder ->
//                    builder.query(esqlQuery).format(format)
//            );
//
//            EsqlQueryResponse response = elasticsearchClient.esql().query(request);
//
//            Map<String, Object> result = new HashMap<>();
//            result.put("operation", "esql");
//            result.put("success", true);
//            result.put("timestamp", Instant.now().toString());
//            result.put("query", esqlQuery);
//            result.put("format", format);
//
//            // 处理不同格式的响应
//            if ("json".equals(format)) {
//                result.put("columns", response.columns());
//                result.put("values", response.values());
//                result.put("total_rows", response.values() != null ? response.values().size() : 0);
//            } else {
//                // CSV或TSV格式，返回原始响应
//                result.put("raw_response", response.toString());
//            }
//
//            return objectMapper.writeValueAsString(result);
//
//        } catch (Exception e) {
//            throw new RuntimeException("执行ES|QL查询失败: " + e.getMessage(), e);
//        }
//    }
//
//    private String getShards(String indexPattern) throws Exception {
//        try {
//            // 使用cat shards API获取分片信息
//            var request = co.elastic.clients.elasticsearch.cat.ShardsRequest.of(builder -> {
//                if (indexPattern != null && !indexPattern.isEmpty()) {
//                    builder.index(indexPattern);
//                }
//                return builder.format("json")
//                        .h("index", "shard", "prirep", "state", "docs", "store", "node");
//            });
//
//            var response = elasticsearchClient.cat().shards(request);
//
//            List<Map<String, Object>> shards = new ArrayList<>();
//            for (var record : response.valueBody()) {
//                Map<String, Object> shardInfo = new HashMap<>();
//                shardInfo.put("index", record.index());
//                shardInfo.put("shard", record.shard());
//                shardInfo.put("prirep", record.prirep());
//                shardInfo.put("state", record.state());
//                shardInfo.put("docs", record.docs());
//                shardInfo.put("store", record.store());
//                shardInfo.put("node", record.node());
//                shards.add(shardInfo);
//            }
//
//            Map<String, Object> result = new HashMap<>();
//            result.put("operation", "get_shards");
//            result.put("success", true);
//            result.put("timestamp", Instant.now().toString());
//            result.put("index_pattern", indexPattern != null ? indexPattern : "*");
//            result.put("total_shards", shards.size());
//            result.put("shards", shards);
//
//            return objectMapper.writeValueAsString(result);
//
//        } catch (Exception e) {
//            throw new RuntimeException("获取分片信息失败: " + e.getMessage(), e);
//        }
//    }
//
//    private String createErrorResponse(String message, String errorType) {
//        Map<String, Object> error = new HashMap<>();
//        error.put("success", false);
//        error.put("error", true);
//        error.put("error_type", errorType);
//        error.put("message", message);
//        error.put("timestamp", Instant.now().toString());
//
//        try {
//            return objectMapper.writeValueAsString(error);
//        } catch (Exception e) {
//            // 如果连序列化都失败，返回硬编码的错误
//            return String.format(
//                    "{\"success\":false,\"error\":true,\"error_type\":\"SerializationError\",\"message\":\"Failed to serialize error: %s\",\"timestamp\":\"%s\"}",
//                    message,
//                    Instant.now().toString()
//            );
//        }
//    }
//}
//.inputSchema("""
//                    {
//                      "type": "object",
//                      "properties": {
//                        "operation": {
//                          "type": "string",
//                          "enum": ["list_indices", "get_mappings", "search", "esql"],
//                          "description": "要执行的操作类型"
//                        },
//                        "index": {
//                          "type": "string",
//                          "description": "索引名称，用于get_mappings、search和get_shards操作"
//                        },
//                        "query": {
//                          "type": "object",
//                          "description": "Elasticsearch查询DSL，用于search操作"
//                        },
//                        "esql_query": {
//                          "type": "string",
//                          "description": "ES|QL查询语句"
//                        },
//                        "size": {
//                          "type": "integer",
//                          "default": 10,
//                          "maximum": 100,
//                          "description": "结果偏移量，用于分页"
//                        }
//                      },
//                      "required": ["operation"]
//                    }
//                    """)
